import logging
import re
import json
import asyncio
from typing import List, Dict, Optional, Tuple, Any

import chat
from chat import split_text, call_llm_api

# --- 配置 ---
logging.basicConfig(format="%(asctime)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S",
                    level=logging.INFO)
logger = logging.getLogger(__name__)


# --- 通用JSON抽取函数 ---
def extract_json_from_response(response: str) -> Optional[object]:
    """
    从模型返回的文本中提取并解析JSON（支持对象和数组）。
    自动清理markdown代码块、多余文字，避免解析失败。
    """
    if not response:
        return None

    # 清理markdown代码块
    response = re.sub(r"```.*?```", "", response, flags=re.DOTALL).strip()

    # 尝试直接解析
    try:
        return json.loads(response)
    except Exception:
        pass

    # 如果失败，尝试用正则抽取第一个JSON对象或数组
    json_match = re.search(r'(\{.*}|\[.*])', response, re.DOTALL)
    if not json_match:
        return None

    candidate = json_match.group(0)
    try:
        return json.loads(candidate)
    except Exception as e:
        logger.error(f"JSON解析失败: {e}")
        return None


# --- 辅助函数 (可共用) ---
def call_chat_api_non_stream(messages: List[Dict]) -> Optional[str]:
    """
    一个辅助函数，用于以非流式方式调用chat.py中的问答API。
    """
    try:
        response = chat.call_llm_api(messages, stream=False)
        if response and isinstance(response, dict):
            return response.get("choices", [{}])[0].get("message", {}).get("content", "")
        return None
    except Exception as e:
        logger.error(f"执行非流式API调用时发生错误: {e}", exc_info=True)
        return None


# --- 功能一：通用实体类型关系提取  ---
def build_unified_schema_extraction_prompt(text: str) -> str:
    """
    通过精确定义和示例，引导模型自主提取正确的实体类型和关系。
    特别强调只输出纯JSON格式，不包含任何额外内容。
    """
    return f"""# 指令
你是一个精确的JSON数据生成器。你的唯一任务是从文本中提取实体类型和关系，并以严格的JSON格式输出结果。

## 任务要求
1.  **实体类型 (Entity Types)**: 指一类具体或抽象的**事物**的类别，其名称通常是**名词或名词性短语**。
    -   **合格的例子**: `人物`, `组织机构`, `地点`, `城市`, `国家`, `日期`, `时间`, `事件`, `战争`, `产品`, `软件`, `书籍`, `电影`, `法律`, `政策`, `疾病`, `药物`。
    -   **绝对不合格的例子 (严禁抽取)**:
        -   **动词/动作类**: `实现`, `评估`, `监督`, `决策`, `感知`, `推进`, `驱动`, `改革`, `发展`。
        -   **抽象领域/概念类**: `时期`, `现代化`, `网络安全`, `工业化`, `城镇化`, `农业现代化`, `社会化处理`, `科技创新`, `集约化管理`, `业务协同`, `安全防护`, `水利信息化`等。
        -   **形容词/状态类**: `智能化`, `大型`, `高效`。
        -   **具体实体名称**: `水利信息基础设施`, `阿里公司`, `北京`, `科技发展大会`, `云服务`等。

2.  **关系 (Relations)**: 连接两个实体类型的**动作**或**语义联系**，其名称通常是**动词或动词性短语**。
    -   **合格的例子**: `出生于`, `颁布了`, `导致了`, `成立于`, `投资了`, `治疗`, `写于`, `是首都`, `隶属于`。
    -   **不合格的例子**: `是`, `有`, `包括`, `称为` (这些关系信息量过低，必须避免)。

## 输出格式要求
你必须输出一个严格的JSON对象，且只包含以下两个字段：
{{
  "entity_types": ["人物", "组织机构", "地点", "政策"],
  "relationships": [
    {{"source": "人物", "relation": "出生于", "target": "地点"}},
    {{"source": "组织", "relation": "颁布了", "target": "政策"}}
  ]
}}

## 绝对规则
1.  你的输出必须是100%纯JSON，不能有任何额外的文本、解释、思考过程或代码块标记。
2.  输出必须直接以 `{{` 开始，以 `}}` 结束。
3.  所有字符串必须使用双引号，不能使用单引号。
4.  关系中的 `source` 和 `target` 必须来自你提取的 `entity_types` 列表。
5.  如果找不到任何实体或关系，返回空数组但仍保持JSON格式。

## 待分析的文本
---
{text}
---
"""

def build_relationship_construction_prompt_with_fixed_types(text: str, entity_types: List[str]) -> str:
    entity_list_str = ", ".join(f'"{t}"' for t in entity_types)
    return f"""# 指令
你是一个精确的JSON数据生成器。你的唯一任务是根据"允许的实体类型列表"在根据提供的文本构建关系，并以严格的JSON数组格式输出结果。

## 任务要求
1.  你必须严格遵守"允许的实体类型列表"，只能使用列表中的实体类型。
2.  如果文本中的关系涉及列表之外的实体类型，必须完全忽略该关系，不得提取。
3.  每个关系必须包含三个字段："source"、"relation"、"target"。
4.  关系必须是文本中明确或隐含描述的，不能凭空创造。
5.  如果提供的允许实体类型与文中的其他实体类型存在关系，不要输出
6.  如果提供的实体类型某两个之间不存在关系，则没必要输出
## 输出格式要求
你必须输出一个严格的JSON数组，格式如下：
[
  {{"source": "实体类型1", "relation": "关系动词", "target": "实体类型2"}},
  {{"source": "实体类型3", "relation": "关系动词", "target": "实体类型4"}}
]

## 绝对规则
1.  你的输出必须是100%纯JSON，不能有任何额外的文本、解释、思考过程或代码块标记。
2.  输出必须直接以 `[` 开始，以 `]` 结束。
3.  所有字符串必须使用双引号，不能使用单引号。
4.  "source" 和 "target" 必须严格来自允许的实体类型列表，不能使用列表外的任何类型。
5.  如果找不到任何符合条件的关系，返回空数组 `[]`。
6.  严禁使用列表外的实体类型，即使它们在文本中出现。

## 允许的实体类型列表 (必须严格使用以下词汇)
[{entity_list_str}]

## 示例
假设允许的实体类型列表是: ["人物", "组织", "地点"]
文本内容是: "张三在阿里巴巴工作，该公司位于杭州。"

正确输出:
[
  {{"source": "人物", "relation": "工作于", "target": "组织"}},
  {{"source": "组织", "relation": "位于", "target": "地点"}}
]

错误输出 (包含不允许的实体类型):
[
  {{"source": "人物", "relation": "工作于", "target": "公司"}},  # "公司"不在允许列表中
  {{"source": "公司", "relation": "位于", "target": "城市"}}   # "公司"和"城市"都不在允许列表中
]

## 待分析文本
---
{text}
---
"""


def extract_entity_relationships_optimized(text: str, instruction: Optional[str] = None) -> List[Dict]:
    """
    主处理函数：根据是否存在instruction，智能选择不同的提取策略。
    """
    if not text:
        return []

    relationships = []

    if instruction:
        logger.info("检测到instruction，将根据预设的实体类型提取关系...")
        entity_types = [item.strip() for item in instruction.split(',') if item.strip()]
        if not entity_types:
            logger.warning("instruction中未包含有效的实体类型。")
            return []

        prompt = build_relationship_construction_prompt_with_fixed_types(text, entity_types)
        messages = [{"role": "user", "content": prompt}]
        response = call_chat_api_non_stream(messages)

        if not response:
            return [{"error": "模型调用失败"}]

        # 尝试提取JSON
        try:
            # 尝试直接解析整个响应
            data = json.loads(response)
            if isinstance(data, list):
                relationships = data
            else:
                logger.error(f"响应不是JSON数组: {response}")
                return [{"error": "模型返回的不是有效的JSON数组"}]
        except json.JSONDecodeError:
            # 如果直接解析失败，尝试从响应中提取JSON部分
            json_match = re.search(r'\[.*]', response, re.DOTALL)
            if json_match:
                try:
                    relationships = json.loads(json_match.group())
                except json.JSONDecodeError as e:
                    logger.error(f"无法从响应中提取有效JSON: {e}")
                    return [{"error": f"JSON解析失败: {e}"}]
            else:
                logger.error(f"响应中未找到JSON数组: {response}")
                return [{"error": "响应中未找到有效的JSON数组"}]

    else:
        logger.info("未提供instruction，将由模型统一提取实体类型和关系...")
        prompt = build_unified_schema_extraction_prompt(text)
        messages = [{"role": "user", "content": prompt}]
        response = call_chat_api_non_stream(messages)

        if not response:
            return [{"error": "模型调用失败"}]

        # 尝试提取JSON
        try:
            # 尝试直接解析整个响应
            data = json.loads(response)
            if isinstance(data, dict) and "relationships" in data:
                relationships = data.get("relationships", [])
            else:
                logger.error(f"响应不是包含relationships字段的JSON对象: {response}")
                return [{"error": "模型返回的不是有效的JSON对象"}]
        except json.JSONDecodeError:
            # 如果直接解析失败，尝试从响应中提取JSON部分
            json_match = re.search(r'\{.*}', response, re.DOTALL)
            if json_match:
                try:
                    data = json.loads(json_match.group())
                    if isinstance(data, dict) and "relationships" in data:
                        relationships = data.get("relationships", [])
                    else:
                        logger.error(f"提取的JSON不包含relationships字段: {response}")
                        return [{"error": "提取的JSON不包含relationships字段"}]
                except json.JSONDecodeError as e:
                    logger.error(f"无法从响应中提取有效JSON: {e}")
                    return [{"error": f"JSON解析失败: {e}"}]
            else:
                logger.error(f"响应中未找到JSON对象: {response}")
                return [{"error": "响应中未找到有效的JSON对象"}]

    final_data = []
    if isinstance(relationships, list):
        for item in relationships:
            if isinstance(item, dict) and all(k in item for k in ["source", "relation", "target"]):
                final_data.append({
                    "srcFieldName": item["source"],
                    "relationName": item["relation"],
                    "dstFieldName": item["target"]
                })

    logger.info(f"流程成功：最终提取到 {len(final_data)} 个关系。")
    return final_data

# --- 功能二：根据模板提取实体实例 ---
def _parse_template(template_str: str) -> Tuple[str, str, str]:
    """
    从 '(源类型)-[关系]-(目标类型)' 格式的字符串中解析三元组。
    """
    match = re.match(r'^\s*\((.*?)\)\s*-\s*\[(.*?)\]\s*-\s*\((.*?)\)\s*$', template_str)
    if match:
        return match.group(1), match.group(2), match.group(3)
    parts = template_str.split('-')
    if len(parts) == 3:
        return parts[0].strip(), parts[1].strip(), parts[2].strip()
    return "未知", "未知", "未知"


def parse_and_number_instructions(instruction: str) -> Dict:
    """
    解析 instruction 字符串，为每个模板分配唯一编号。
    """
    templates = [item.strip() for item in instruction.split(',') if item.strip()]
    numbered_list_for_prompt = []
    number_to_template_map = {}
    for i, tpl_str in enumerate(templates, 1):
        numbered_list_for_prompt.append(f"[{i}] {tpl_str}")
        number_to_template_map[str(i)] = tpl_str
    return {"prompt_list": numbered_list_for_prompt, "map": number_to_template_map}


def build_numbered_extraction_prompt(text: str, numbered_templates_batch: List[str], more_type: bool) -> str:
    """
    根据 more_type 的值构建指令，并增加了更严格的提取规则。
    """
    templates_str = "\n".join(numbered_templates_batch)
    if more_type:
        task_description = """你的任务是：根据我提供的“带编号的实体类型对列表”，从文本中找到连接这两种类型实体的所有可能关系，并抽取出具体的三元组实例。"""
        example_block = """
**示例任务**:
- 文本: "苹果公司于昨日发布了新款iPhone 18，该产品由其设计团队精心打造。同时，微软也宣布收购了一家AI创业公司。"
- 实体类型对列表: [1] (公司) 和 (产品)
**你的思考过程**:
1.  **识别实体**: 我需要在文本中找到所有“公司”和“产品”类型的实体。公司: "苹果公司", "微软"。产品: "新款iPhone 18"。
2.  **寻找关系**: 我需要寻找连接已识别的“公司”和“产品”实体的动词或短语。"苹果公司" 和 "新款iPhone 18" 之间有关联，文本中写道“发布了”。文本还提到产品“由其设计团队精心打造”，这里的“其”指代“苹果公司”，所以“打造”也是一个关系。
3.  **构建输出**: 根据找到的关系，为编号 "1" 构建JSON数组。`{"source": "苹果公司", "relation": "发布了", "target": "新款iPhone 18"}`, `{"source": "苹果公司", "relation": "打造", "target": "新款iPhone 18"}`
**最终示例输出**: `{{"1": [{{"source": "苹果公司", "relation": "发布了", "target": "新款iPhone 18"}}, {{"source": "苹果公司", "relation": "打造", "target": "新款iPhone 18"}}]}}`
"""
        prompt_header = f"""你是一个精准的信息抽取引擎。{task_description}\n{example_block}\n**你的任务正式开始**:\n**带编号的实体类型对列表**:\n{templates_str}"""
    else:
        task_description = """你的任务是根据我提供的“带编号的关系模板列表”，从文本中找到所有匹配的**具体实体实例**。\n**重要**：请**严格按照**模板中指定的关系（例如“-召开-”）进行提取。"""
        prompt_header = f"""你是一个精准的信息抽取引擎。{task_description}\n**带编号的关系模板列表**:\n{templates_str}"""

    # 在规则中增加了第 3 和第 4 条，以提升提取质量
    return f"""{prompt_header}\n\n请严格遵循以下规则：
    1.  **为每个编号提取**：你必须为列表中的**每一个**编号进行独立的提取。
    2.  **利用描述**：请利用括号内的类型描述来准确识别实体。
    3.  **严格对应类型**：提取的`source`实体必须符合模板中的源实体类型，`target`实体必须符合目标实体类型。
    4.  **关系词要求**：提取的`relation`必须是描述动作的**核心动词**或**简短动词短语**（例如：“签署”、“属于”、“投资了”）。它不应是名词或长句子。
    5.  **输出格式**：你的回复必须是一个严格的JSON字符串，代表一个**字典(object)**。
        -   **键 (key)**：是对应的**编号** (字符串格式，例如: "1", "2")。
        -   **值 (value)**：是一个**数组**，包含所有找到的三元组。每个三元组是一个包含`"source"`, `"relation"`, `"target"`键的字典。
        -   如果某个编号在文本中找不到任何实例，其对应的值应该是一个空数组 `[]`。
    6.  **【至关重要的要求】**：
        -   你的回复**必须**只能包含JSON文本。
        -   **绝对不要**添加任何解释、说明或代码块标记。
        -   你的回复必须直接以 '{{' 字符开始，并以 '}}' 字符结束。
    7.  **分片文本意识**：你处理的文本可能只是一个片段。如果当前片段中没有找到任何符合模板的实例，返回一个空数组 `[]` 是完全正确的行为。
\n\n现在，请处理以下文本内容：\n---\n{text}\n---"""


async def _extract_for_batch_async(text_chunk: str, numbered_templates_batch: List[str], more_type: bool) -> Dict:
    """
    [内部函数] 异步处理单个批次、单个文本块的任务。失败时抛出异常。
    """
    if not numbered_templates_batch: return {}
    batch_keys = [re.match(r'\[(\d+)\]', tpl).group(1) for tpl in numbered_templates_batch if
                  re.match(r'\[(\d+)\]', tpl)]

    if more_type:
        processed_batch_for_prompt = []
        for tpl in numbered_templates_batch:
            num = re.match(r'\[(\d+)\]', tpl).group(1)
            template_content = re.sub(r'\[\d+\]\s*', '', tpl)
            source_type, _, target_type = _parse_template(template_content)
            processed_batch_for_prompt.append(f"[{num}] 源实体类型({source_type}) -> 目标实体类型({target_type})")
        prompt = build_numbered_extraction_prompt(text_chunk, processed_batch_for_prompt, more_type)
    else:
        prompt = build_numbered_extraction_prompt(text_chunk, numbered_templates_batch, more_type)

    messages = [{"role": "user", "content": prompt}]

    try:
        loop = asyncio.get_running_loop()
        model_response_dict = await loop.run_in_executor(None, call_llm_api, messages, False)

        if not model_response_dict:
            raise IOError("模型未能返回有效内容")

        model_response_text = model_response_dict.get("choices", [{}])[0].get("message", {}).get("content", "")
        if not model_response_text:
            return {}

        json_match = re.search(r'\{.*\}', model_response_text, re.DOTALL)
        if not json_match:
            raise ValueError("模型返回内容不是有效的JSON结构")

        return json.loads(json_match.group(0))

    except Exception as e:
        logger.error(f"处理批次 {batch_keys} 时发生严重错误: {e}")
        raise RuntimeError(f"处理文本块时失败: {e}") from e


async def extract_specific_entities_concurrently(text: str, instruction: str, more_type: bool = False) -> Dict[
    str, Any]:
    """
    [主处理函数] 提取实体实例，采用新的动态分组嵌套逻辑。
    """
    try:
        if not text or not instruction:
            return {"code": 400, "msg": "请求参数'text'或'instruction'不能为空", "data": None}

        # 1. 预处理
        parsed_data = parse_and_number_instructions(instruction)
        prompt_list = parsed_data["prompt_list"]
        number_to_template_map = parsed_data["map"]

        if not prompt_list:
            return {"code": 400, "msg": "从'instruction'中未能解析出任何有效的提取模板", "data": None}

        # 2. 分片与任务创建
        text_chunks = split_text(text, chunk_size=20000, overlap=1000)
        batch_size = 10
        tasks = []

        for chunk in text_chunks:
            batches = [prompt_list[i:i + batch_size] for i in range(0, len(prompt_list), batch_size)]
            for batch in batches:
                tasks.append(_extract_for_batch_async(chunk, batch, more_type))

        # 3. 并发执行
        logger.info(f"为 {len(text_chunks)} 个文本块创建了 {len(tasks)} 个并发任务，开始执行...")
        all_results = await asyncio.gather(*tasks)
        logger.info("所有并发任务已成功完成，开始合并结果...")

        # 4. 合并结果
        combined_results = {}
        for result_dict in all_results:
            for num, instances in result_dict.items():
                if num not in combined_results:
                    combined_results[num] = []
                combined_results[num].extend(instances)

        # 5. 格式化最终输出，实现新的动态分组嵌套逻辑

        # 使用一个字典来按新的 relationship_type 进行分组
        # 键: "源类型-关系-目标类型", 值: [实例列表]
        grouped_results = {}
        seen_instances = set()  # 用于去重完全相同的三元组实例

        for num, instances in combined_results.items():
            if not isinstance(instances, list):
                continue

            rel_type_template = number_to_template_map.get(num)
            if not rel_type_template:
                continue
            source_type, original_relation, target_type = _parse_template(rel_type_template)

            for item in instances:
                if isinstance(item, dict) and all(k in item for k in ["source", "target"]):

                    if more_type:
                        relation_word = item.get("relation", "未知关系")
                    else:
                        relation_word = original_relation

                    # 创建用于去重的唯一标识
                    instance_tuple = (item.get("source"), relation_word, item.get("target"))

                    if instance_tuple not in seen_instances:
                        # 动态构建分组的键 (即新的 relationship_type)
                        group_key = f"{source_type}-{relation_word}-{target_type}"

                        # 格式化单个实例
                        formatted_instance = {
                            "srcFieldName": item["source"],
                            "relationName": relation_word,
                            "dstFieldName": item["target"]
                        }

                        # 将格式化后的实例添加到对应的分组中
                        if group_key not in grouped_results:
                            grouped_results[group_key] = []
                        grouped_results[group_key].append(formatted_instance)

                        seen_instances.add(instance_tuple)

        # 将分组后的字典转换为最终的列表格式
        final_data = []
        for rel_type, inst_list in grouped_results.items():
            final_data.append({
                "relationship_type": rel_type,
                "instances": inst_list
            })

        return {
            "code": 200,
            "msg": "操作成功",
            "data": final_data
        }

    except Exception as e:
        logger.error(f"提取实体实例时发生严重错误: {e}", exc_info=True)
        return {
            "code": 500,
            "msg": f"{e}",
            "data": None
        }


# --- 新增功能一：提取带描述的实体类型 (type_des) ---
def build_type_description_prompt(text: str, user_input: Optional[str] = None) -> str:
    """
    构建一个指令，用于提取实体类型及其在文中的描述。
    """
    instruction_section = f"用户的额外要求（请优先满足这些要求）：{user_input}" if user_input else "请全面地识别所有实体类型及其描述。"

    return f"""你是一个高级的概念定义抽取专家。你的任务是从我提供的文本中，识别出关键的**实体类型**，并对该实体类型生成非文本内容的**具体描述**。

请严格遵循以下规则：

1.  **识别实体类型**：识别文本中的通用类别，如“公司”、“会议”、“文件”、“组织”、“时间”、“货币”、“地区”、“城市”、“人员”等。
2.  **生成描述**：根据你自己的对单词的理解，生成对实体类型的通用解释，而不是根据文本内容生成的描述。
3.  **遵守用户要求**：{instruction_section}
4.  **输出格式**：你的回复必须是一个严格的JSON字符串，代表一个对象数组。每个对象必须包含两个键：
    -   `"entity_type"`: (字符串) 识别出的实体类型。
    -   `"description"`: (字符串) 对每个实体类型的对应描述。
5.  **【至关重要的要求】**：
    -   如果某个实体类型在文中没有明确的描述，请不要包含在结果中。
    -   你的回复必须只能包含JSON文本，不要包含任何解释或代码块标记。
    -   你的回复必须直接以 `[` 字符开始，并以 `]` 字符结束。
    -   示例输出: `[{{"entity_type": "公司", "description": "依法设立的盈利性商业组织实体。"}}]`

现在，请处理以下文本内容：

---
{text}
---
"""


async def _extract_types_from_chunk_async(text_chunk: str, user_input: Optional[str] = None) -> List[Dict]:
    """
    [内部函数] 异步处理单个文本块的实体类型提取任务。
    这是并发处理中的 "Map" 步骤。
    """
    logger.info(f"开始处理一个文本块，长度为 {len(text_chunk)}...")
    prompt = build_type_description_prompt(text_chunk, user_input)
    messages = [{"role": "user", "content": prompt}]

    try:
        loop = asyncio.get_running_loop()
        # 使用 to_thread 在后台运行同步的API调用，并设置 stream=False 获取完整JSON
        model_response_dict = await loop.run_in_executor(
            None,
            call_llm_api,
            messages, False  # stream=False
        )

        if not model_response_dict:
            logger.warning("一个文本块的模型调用未能返回有效内容。")
            return []

        # 非流式调用直接返回字典，提取其中的内容
        model_response_text = model_response_dict.get("choices", [{}])[0].get("message", {}).get("content", "")

        json_match = re.search(r'\[.*]', model_response_text, re.DOTALL)
        if not json_match:
            raise ValueError("在模型返回的文本中未找到JSON数组结构")

        results = json.loads(json_match.group(0))
        if not isinstance(results, list):
            raise TypeError("解析出的JSON不是列表")

        logger.info(f"成功从文本块中提取到 {len(results)} 个实体类型。")
        return results

    except Exception as e:
        logger.error(f"处理文本块时发生错误: {e}")
        # 返回一个包含错误信息的特殊对象，以便上层函数识别
        return [{"internal_error": f"处理文本块失败: {str(e)}"}]


async def extract_entity_types_with_description(text: str, user_input: Optional[str] = None) -> Dict[str, Any]:
    """
    [主处理函数] 调用模型提取实体类型及其描述，支持长文本，并按要求格式化输出。
    """
    try:
        if not text:
            # 兼容旧逻辑，但建议在API层处理空文本
            return {"code": 200, "msg": "操作成功", "data": []}

        # 1. 分片 (Split) - 复用 chat.py 的逻辑
        # 您可以根据模型的上下文窗口调整 chunk_size
        text_chunks = split_text(text, chunk_size=25000, overlap=1000)

        # 2. 并行处理 (Map)
        tasks = [_extract_types_from_chunk_async(chunk, user_input) for chunk in text_chunks]
        chunk_results = await asyncio.gather(*tasks)

        # 3. 合并与去重 (Reduce)
        all_entities = {}
        for result_list in chunk_results:
            # 检查并跳过处理失败的块
            if result_list and isinstance(result_list[0], dict) and "internal_error" in result_list[0]:
                logger.warning(f"跳过一个处理失败的文本块: {result_list[0]['internal_error']}")
                continue

            for entity in result_list:
                if isinstance(entity, dict) and 'entity_type' in entity:
                    entity_type = entity['entity_type']
                    if entity_type not in all_entities:
                        all_entities[entity_type] = entity

        final_results = list(all_entities.values())
        logger.info(f"合并去重后，共得到 {len(final_results)} 个唯一的实体类型。")

        return {
            "code": 200,
            "msg": "操作成功",
            "data": final_results
        }

    except Exception as e:
        logger.error(f"提取实体类型描述时发生严重错误: {e}", exc_info=True)
        return {
            "code": 500,
            "msg": f"服务器内部错误: {str(e)}",
            "data": None
        }


# --- 新增功能二：根据带描述的实体类型构建关系 (type_rel) ---
def unify_and_parse_inputs_with_numbering(instruction: Optional[str], user_input: Optional[str]) -> Dict:
    """
    统一解析 instruction 和 user_input，为每个实体分配编号。
    """
    full_instruction_str = (instruction or "") + ", " + (user_input or "")
    numbered_list_for_prompt = []
    number_to_full_map = {}
    parts = [p.strip() for p in full_instruction_str.split(',') if p.strip()]
    i = 1
    for part in parts:
        prompt_entry = f"[{i}] {part}"
        numbered_list_for_prompt.append(prompt_entry)
        number_to_full_map[str(i)] = part
        i += 1
    return {"prompt_list": numbered_list_for_prompt, "map": number_to_full_map}


def build_numbered_relationship_prompt(text: Optional[str], numbered_entities: List[str],
                                       user_input: Optional[str] = None) -> str:
    """
    根据 text 是否存在，构建不同的指令。
    """
    instruction_section = f"用户的额外要求（请优先满足这些要求）：{user_input}" if user_input else "请构建所有在文本中明确存在的、或基于常识最可能存在的有意义的关系。"
    entity_list_str = "\n".join(numbered_entities)

    if text:
        task_description = f"""你是一个高级的语义关系分析专家。你的任务是从我提供的**文本**和“带编号的实体列表”中，构建这些实体之间的关系。

**带编号的实体列表**:
---
{entity_list_str}
---

请严格遵循以下规则：

1.  **限定范围**：你构建关系时，必须使用实体对应的**编号**。
2.  **推断关系**：**分析文本**，找到连接这些实体实例的通用动词或上下文，形成关系。
3.  **遵守用户要求**：{instruction_section},用户要求新增或者限定关系时，应优先按照用户的要求进行处理

现在，请处理以下文本内容：
---
{text}
---
"""
    else:
        task_description = f"""你是一个高度智能的知识图谱专家。你的任务是**仅根据**我提供的“带编号的实体列表”中实体的**字面意思和你的常识**，来推断它们之间最可能的通用关系。

**带编号的实体列表**:
---
{entity_list_str}
---

请严格遵循以下规则：

1.  **限定范围**：你构建关系时，必须使用实体对应的**编号**。
2.  **推断关系**：**基于常识**推断实体间的关系。例如，"公司"很可能"召开" "会议"，"主持人"很可能"主持" "会议"。
3.  **遵守用户要求**：{instruction_section}
"""

    return task_description + f"""
**输出格式要求 (对所有情况都适用)**:
-   你的回复必须是纯文本，每行一个关系，并严格遵循 `[源实体编号] -> 关系 -> [目标实体编号]` 的格式。
-   **绝对不要添加任何解释、说明、编号或代码块标记**。
-   你的回复必须直接以第一个关系三元组开始。
-   如果找不到任何关系，请不要输出任何内容。
-   **示例输出**:
[1] -> 召开 -> [2]
[3] -> 主持 -> [2]
"""


async def _extract_relationships_from_chunk_async(
        text_chunk: Optional[str],
        numbered_entities: List[str],
        user_input: Optional[str]
) -> str:
    """
    [内部函数] 异步处理单个文本块的关系提取任务。返回原始文本结果。
    """
    prompt = build_numbered_relationship_prompt(text_chunk, numbered_entities, user_input)
    messages = [{"role": "user", "content": prompt}]

    try:
        loop = asyncio.get_running_loop()
        # 使用 to_thread 在后台运行同步的API调用，并设置 stream=False
        model_response_dict = await loop.run_in_executor(
            None,
            call_llm_api,
            messages, False  # stream=False
        )

        if not model_response_dict:
            logger.warning("一个文本块的模型调用未能返回有效内容。")
            return ""  # 返回空字符串，避免影响后续处理

        # 从返回的字典中提取纯文本内容
        return model_response_dict.get("choices", [{}])[0].get("message", {}).get("content", "")

    except Exception as e:
        logger.error(f"处理文本块关系提取时发生错误: {e}")
        return ""  # 同样返回空字符串


async def build_relationships_with_description(
        text: str,
        instruction: str,
        user_input: Optional[str] = None
) -> Dict[str, Any]:
    """
    [主处理函数-修改版] 智能处理有无文本以及长文本情况，并按实体对进行去重。
    """
    try:
        # 1. 预处理：解析输入并为实体编号
        parsed_data = unify_and_parse_inputs_with_numbering(instruction, user_input)
        prompt_list = parsed_data["prompt_list"]
        number_to_full_map = parsed_data["map"]

        if not prompt_list:
            logger.warning("从输入中未能解析出任何有效的实体。")
            return {"code": 200, "msg": "操作成功", "data": []}

        # 2. 调用模型获取关系 (区分有无文本场景)
        model_response = ""
        if text and text.strip():
            # 场景一：有文本，启动长文本处理流程
            logger.info("检测到输入文本，启动长文本关系提取流程...")
            text_chunks = split_text(text, chunk_size=25000, overlap=1000)
            tasks = [_extract_relationships_from_chunk_async(chunk, prompt_list, user_input) for chunk in text_chunks]
            chunk_results = await asyncio.gather(*tasks)
            model_response = "\n".join(chunk_results)
        else:
            # 场景二：无文本，执行单次常识推理
            logger.info("未检测到输入文本，启动基于常识的关系推理...")
            model_response = await _extract_relationships_from_chunk_async(None, prompt_list, user_input)

        # 3. 后处理：解析模型的文本输出并重组
        final_data = []
        # 这里的 seen_relations 现在将存储已见过的实体对
        seen_entity_pairs = set()  # 用于按实体对去重
        pattern = re.compile(r"\[(\d+)]\s*->\s*(.+?)\s*->\s*\[(\d+)]")
        matches = pattern.finditer(model_response)

        for match in matches:
            source_num, relation, target_num = match.groups()

            # 创建不区分方向的唯一实体对键
            # 通过排序，确保 (1, 2) 和 (2, 1) 会生成相同的键
            pair_key = tuple(sorted((source_num, target_num)))

            # 检查实体编号是否存在，并且这个实体对是第一次出现
            if source_num in number_to_full_map and target_num in number_to_full_map and pair_key not in seen_entity_pairs:
                final_data.append({
                    "source": number_to_full_map[source_num],
                    "relation": relation.strip(),
                    "target": number_to_full_map[target_num]
                })
                # 将这个实体对标记为已处理
                seen_entity_pairs.add(pair_key)

        logger.info(f"成功从模型响应中解析并重组了 {len(final_data)} 个唯一关系（按实体对去重）。")
        return {
            "code": 200,
            "msg": "操作成功",
            "data": final_data
        }

    except Exception as e:
        logger.error(f"构建实体关系时发生严重错误: {e}", exc_info=True)
        return {
            "code": 500,
            "msg": f"服务器内部错误: {str(e)}",
            "data": None
        }


# --- 新增功能三：实体消歧 (entity_dis) ---
def build_disambiguation_prompt(user_input: str) -> str:
    """
    构建一个指令，用于对输入的文本/JSON进行实体消歧。
    """
    return f"""你是一个顶级的实体解析和指代消解专家。你的任务是分析我提供的输入内容，识别出所有指向同一真实世界对象的同义实体，并将它们统一为一个规范化的名称。

请严格遵循以下规则：

1.  **识别同义词**：找出所有意义相同但名称略有不同的实体。例如，“北京”和“北京市”应被视为同义词。
2.  **选择规范名称**：从同义词中选择一个最常用或最完整的作为“规范名称”。
3.  **创建映射**：生成一个JSON对象，其中“键”是原始的、需要被替换的同义词，“值”是统一后的规范名称。
4.  **【至关重要的要求】**：
    -   你的回复必须是一个严格的JSON字符串，代表一个**字典(object)**。
    -   **只包含需要统一的实体**。如果一个实体没有同义词，请不要包含在结果中。
    -   你的回复必须只能包含JSON文本，不要包含任何解释或代码块标记。
    -   你的回复必须直接以 `{{` 字符开始，并以 `}}` 字符结束。
    -   示例输出: `{{"北京市": "北京", "B.J.": "北京"}}`

现在，请分析以下输入内容，并生成同义实体的映射：

---
{user_input}
---
"""


def disambiguate_entities(user_input: str) -> str:
    """
    主处理函数：调用模型获取实体映射，然后在Python中执行替换，返回处理后的原始文本。
    """
    if not user_input:
        return ""

    prompt = build_disambiguation_prompt(user_input)
    messages = [{"role": "user", "content": prompt}]

    logger.info("正在调用模型进行实体消歧...")
    model_response = call_chat_api_non_stream(messages)

    if not model_response:
        logger.error("模型未能返回有效的消歧映射。")
        return user_input  # 出错时返回原文

    try:
        json_match = re.search(r'\{.*}', model_response, re.DOTALL)
        if not json_match: raise ValueError("在模型返回的文本中未找到JSON对象结构")

        disambiguation_map = json.loads(json_match.group(0))
        if not isinstance(disambiguation_map, dict): raise TypeError("解析出的JSON不是字典")

        logger.info(f"成功获取消歧映射: {disambiguation_map}")

        # 在Python中执行替换，这比让模型直接修改更可靠
        processed_input = user_input
        for original, replacement in disambiguation_map.items():
            # 使用 re.sub 进行全词匹配替换，避免替换掉单词的一部分
            processed_input = re.sub(r'\b' + re.escape(original) + r'\b', replacement, processed_input)

        return processed_input

    except (json.JSONDecodeError, ValueError, TypeError) as e:
        logger.error(f"解析消歧映射JSON时失败: {e}", extra={'raw_model_output': model_response})
        return user_input  # 出错时返回原文
